# 使用随机森林完成泰坦尼克号幸存者预测。

import pandas as pd
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import cross_val_score
from sklearn.ensemble import RandomForestClassifier
import matplotlib.pyplot as plt

# 读取数据
data = pd.read_csv('./data/titanic.csv')
print(data.describe())
print(data.head())
# 数据处理

# 处理缺失值，对缺失值较多的列进行填补，有一些特征只缺失一两个值，可以采取直接删除记录的方法
data["age"] = data["age"].fillna(data["age"].mean())
data = data.dropna()

data["sex"] = (data["sex"] == "male").astype("int")
data["adult_male"] = (data["adult_male"] == "TRUE").astype("int")

labels = data["embarked"].unique().tolist()
data["embarked"] = data["embarked"].apply(lambda x: labels.index(x))
labels1 = data["who"].unique().tolist()
data["who"] = data["who"].apply(lambda x: labels1.index(x))
data.drop(columns=['class', 'deck', 'embark_town', 'alive'], inplace=True)

print(data.head())

X = data.iloc[:, data.columns != "survived"]
y = data.iloc[:, data.columns == "survived"]

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)

for i in [X_train, X_test, y_train, y_test]:
    i.index = range(i.shape[0])

clf = RandomForestClassifier(n_estimators=4, max_depth=3, criterion="gini", max_features=0.1, min_samples_split=5)
y_train = y_train.values.ravel()
clf = clf.fit(X_train, y_train)
score_ = clf.score(X_test, y_test)
print(score_)

train_list = []
test_list = []
for i in range(10):
    clf = DecisionTreeClassifier(random_state=4, max_depth=i+1, criterion="entropy")
    clf = clf.fit(X_train, y_train)
    score_train = clf.score(X_train, y_train)
    score_test = cross_val_score(clf, X, y, cv=10).mean()
    train_list.append(score_train)
    test_list.append(score_test)
print(max(test_list))
plt.figure(figsize=(10, 6))
plt.plot(range(1, 11), train_list, color="orange", label="train")
plt.plot(range(1, 11), test_list, color="black", label="test")
plt.xticks(range(1, 11))
plt.legend()
plt.savefig("随机森林.jpg")
plt.show()